"""
Basic skeleton of a mitmproxy addon.

Run as follows: mitmproxy -s anatomy.py
"""

import logging

import yaml

# 定义要拦截的域名
BLOCKED_DOMAIN = "test-pfv1.rongyi.com"

class Counter:
    def __init__(self):
        self.num = 0

    def request(self, flow):
        self.num = self.num + 1
        logging.info("We've seen %d flows" % self.num)
        logging.info(f"Intercepted request: {flow.request.url}")


    def response(self, flow):
        logging.info(f"Response: {flow.response.status_code}")
        logging.info(f"Headers: {flow.response.headers}")
        if BLOCKED_DOMAIN in flow.request.url:
            logging.info(f"Blocking request to {BLOCKED_DOMAIN}: {flow.request.url}")
            logging.info(f"Request: {flow.request.url} {flow.request.method}")
            logging.info(f"Headers: {flow.request.headers}")

            # 响应可能尚未生成，检查是否存在
            if flow.response:
                logging.info(f"Response: {flow.response.status_code}")
                logging.info(f"Headers: {flow.response.headers}")

                # 将请求和响应数据保存到字典中
                flow_data = {
                    "request": {
                        "method": flow.request.method,
                        "url": flow.request.url,
                        "headers": dict(flow.request.headers),
                        "body": flow.request.text if flow.request.text else "",
                    },
                    "response": {
                        "status_code": flow.response.status_code,
                        "headers": dict(flow.response.headers),
                        "body": flow.response.text if flow.response.text else "",
                    }
                }

                # 将数据保存到 YAML 文件中
                with open(f"{flow.request.url.replace('://', '_').replace('/', '_')}.yaml", "w") as file:
                    yaml.dump(flow_data, file)


addons = [Counter()]

# mitmdump -s -p 8866 ./anatomy.py
